import torch as th
import torch.nn as nn
import torch.nn.functional as F
import warnings
warnings.filterwarnings("ignore", category=UserWarning)

class Critic(nn.Module):
    def __init__(self, n_agents, dim_observation, dim_action):
        super(Critic, self).__init__()
        self.n_agents = n_agents
        self.dim_observation = dim_observation
        self.dim_action = dim_action
        obs_dim = dim_observation * n_agents
        act_dim = self.dim_action * n_agents

        self.FC1 = nn.Linear(obs_dim, 1024)
        nn.init.kaiming_normal_(self.FC1.weight,1)
        #nn.init.constant_(self.FC1.bias,0.0)
        self.FC2 = nn.Linear(1024+act_dim, 512)
        nn.init.kaiming_normal_(self.FC2.weight,1)
        #nn.init.constant_(self.FC2.bias,0.0)
        self.FC3 = nn.Linear(512, 300)
        nn.init.kaiming_normal_(self.FC3.weight,1)
        #nn.init.constant_(self.FC3.bias,0.0)
        self.FC4 = nn.Linear(300, 1)
        nn.init.kaiming_normal_(self.FC4.weight,1)
        #nn.init.constant_(self.FC4.bias,0.0)

    # obs: batch_size * obs_dim
    def forward(self, obs, acts):
        result = F.relu(self.FC1(obs))
        combined = th.cat([result, acts], 1)
        result = F.relu(self.FC2(combined))
        return self.FC4(F.relu(self.FC3(result)))


class Actor(nn.Module):
    def __init__(self, n_enemies, dim_observation, dim_action):
        super(Actor, self).__init__()
        self.n_enemies = n_enemies
        self.dim_observation = dim_observation
        self.dim_action = dim_action
        self.FC1 = nn.Linear(self.dim_observation, 500)
        nn.init.kaiming_normal_(self.FC1.weight,1)
        #nn.init.constant_(self.FC1.bias,0.0)
        self.FC2 = nn.Linear(500, 128)
        nn.init.kaiming_normal_(self.FC2.weight,1)
        #nn.init.constant_(self.FC2.bias,0.0)
        self.FC3 = nn.Linear(128, self.dim_action)
        nn.init.kaiming_normal_(self.FC3.weight,1)
        #nn.init.constant_(self.FC3.bias,0.0)

    # action output between -15 and 15
    def forward(self, obs):
        result = F.relu(self.FC1(obs))
        for item in result[0]:
            if th.isnan(item):
                print('NaN1') 
                break  
        result = F.relu(self.FC2(result))
        for item in result[0]:
            if th.isnan(item):
                print('NaN2') 
                break
        result = F.tanh(self.FC3(result))
        for item in result[0]:
            if th.isnan(item):
                print('NaN3') 
                break
        return result
